from __future__ import print_function

import ctypes
import subprocess
import unittest
from ctypes import CFUNCTYPE, c_double, c_int, c_char_p

import llvmlite.binding as llvm
from llvmlite import ir
from llvmlite.binding import (
    targets,
    ExecutionEngine,
    parse_bitcode,
    create_function_pass_manager,
    FunctionPassManager,
    ModuleRef,
)
from llvmlite.ir import Constant, GlobalVariable

from common.utils import log_and_raise
from log import LOGGER, WORK_DIR, write_unopt_ir, write_opt_ir
from spel4llvm.llvm_types import (
    LLVM_INT_TYPE,
    LLVM_FLOAT_TYPE,
    LLVM_CHAR_TYPE,
    LLVM_STR_TYPE,
    LLVM_BOOL_TYPE,
    LLVM_LIST_CONTAINER_TYPE,
    get_llvm_list_full_type,
    is_llvm_list_type,
    extract_llvm_list_elem_type,
    clist_type_factory,
)
from spel4llvm.value_transfer import (
    str_value_to_llvm_storage,
    numeric_value_to_llvm_globalvar,
)


def prepare_basic_block(fn_type: ir.FunctionType):
    # Create an empty module...
    module = ir.Module(name=__file__)
    module.triple = targets.get_default_triple()
    func = ir.Function(module, fn_type, name=DFT_FUNC_NAME)

    # Now implement the function
    block = func.append_basic_block(name="entry")
    return module, block


# All these initializations are required for code generation!
llvm.initialize()
llvm.initialize_native_target()
llvm.initialize_native_asmprinter()  # yes, even this one

DFT_FUNC_NAME = "run"


class LLVMRunner:
    def __init__(self, need_opt=False):
        self.engine = LLVMRunner.create_execution_engine()
        self.need_opt = need_opt
        self.type_mapping = {
            LLVM_INT_TYPE: c_int,
            LLVM_FLOAT_TYPE: c_double,
            LLVM_STR_TYPE: c_char_p,
            LLVM_BOOL_TYPE: ctypes.c_int8,
            LLVM_CHAR_TYPE: ctypes.c_char,
        }

    @staticmethod
    def create_execution_engine():
        """
        Create an ExecutionEngine suitable for JIT code generation on
        the host CPU.  The engine is reusable for an arbitrary number of
        modules.
        """
        # Create a target machine representing the host
        LOGGER.info(targets.get_default_triple())
        target = llvm.Target.from_default_triple()
        target_machine = target.create_target_machine()
        # And an execution engine with an empty backing module
        backing_mod = llvm.parse_assembly("")
        engine = llvm.create_mcjit_compiler(backing_mod, target_machine)
        LLVMRunner.load_ext(engine)
        return engine

    @staticmethod
    def load_ext(engine: ExecutionEngine):
        with open("%s/extensions/ext.bc" % WORK_DIR, "b+r") as fp:
            engine.add_module(parse_bitcode(fp.read()))

    def compile_ir(self, llvm_ir_code):
        """
        Compile the LLVM IR string with the given engine.
        The compiled module object is returned.
        """
        # Create a LLVM module object from the IR
        mod = llvm.parse_assembly(llvm_ir_code)
        mod.verify()

        # Now add the module and make sure it is ready for execution
        self.engine.add_module(mod)
        self.engine.finalize_object()
        self.engine.run_static_constructors()
        return mod

    def optimize(self, mod: ModuleRef):
        LOGGER.info("llvm_runner.optimize begin")
        fpm: FunctionPassManager = create_function_pass_manager(mod)
        fpm.add_aggressive_dead_code_elimination_pass()
        # fpm.add_constant_merge_pass()
        fpm.add_dead_code_elimination_pass()
        fpm.add_dead_store_elimination_pass()
        fpm.initialize()
        fpm.run(mod.get_function(DFT_FUNC_NAME))
        fpm.finalize()
        LOGGER.info("llvm_runner.optimize end")

    def cvt_bc2ir(self, bitcode):
        tmp_bc_file = "%s/tmp.bc" % WORK_DIR
        with open(tmp_bc_file, mode="b+w") as fp:
            fp.write(bitcode)

        tmp_ir_file = "%s/tmp.ll" % WORK_DIR
        proc = subprocess.run(
            ["llvm-dis", tmp_bc_file, "-o", tmp_ir_file], capture_output=True, text=True
        )
        if proc.returncode == 0:
            with open(tmp_ir_file, mode="r") as fp:
                return fp.read()
        else:
            LOGGER.error("bc2ir failed, returncode:%d" % proc.returncode)
            if proc.stderr:
                LOGGER.error("bc2ir error detail:%s" % proc.stderr)
            return None

    def run(self, ir_module: ir.Module, *args):
        llvm_ir_code = str(ir_module)
        LOGGER.info("llvm ir:%s" % llvm_ir_code)
        mod = self.compile_ir(llvm_ir_code)

        if self.need_opt:
            write_unopt_ir(self.cvt_bc2ir(mod.as_bitcode()))
            self.optimize(mod)
            write_opt_ir(self.cvt_bc2ir(mod.as_bitcode()))

        # Look up the function pointer (a Python int)
        func_ptr = self.engine.get_function_address(DFT_FUNC_NAME)

        func_defs = [x for x in ir_module.functions if x.name == DFT_FUNC_NAME]
        if len(func_defs) == 0:
            raise SyntaxError("no function:" + DFT_FUNC_NAME)

        dft_func_def = func_defs[0]
        return_type = self.get_c_type(dft_func_def.return_value.type)
        arg_types = [self.get_c_type(x.type) for x in dft_func_def.args]

        LOGGER.info(arg_types)
        # Run the function via ctypes
        cfunc = CFUNCTYPE(return_type, *arg_types)(func_ptr)
        return cfunc(*args)

    def get_global_value_address(self, var_name):
        return self.engine.get_global_value_address(var_name)

    def get_c_type(self, ir_type):
        c_tp = self.type_mapping.get(ir_type)
        if c_tp is None:
            # process list type first
            if isinstance(ir_type, LLVM_LIST_CONTAINER_TYPE):
                return self.get_c_type(ir_type.element) * ir_type.count
            elif is_llvm_list_type(ir_type):
                return clist_type_factory(
                    self.get_c_type(extract_llvm_list_elem_type(ir_type))
                )
            # at last process Pointer
            elif isinstance(ir_type, ir.PointerType):
                return ctypes.POINTER(self.get_c_type(ir_type.pointee))
            else:
                log_and_raise(SyntaxError, "cannot find ctype of %s" % ir_type)
        else:
            return c_tp


class LLVMRunnerTest(unittest.TestCase):
    @staticmethod
    def gen_ir_i32_module():
        # Create some useful types
        fnty = ir.FunctionType(LLVM_INT_TYPE, (LLVM_INT_TYPE, LLVM_INT_TYPE))

        # Create an empty module...
        module = ir.Module(name=__file__)
        module.triple = targets.get_default_triple()
        func = ir.Function(module, fnty, name=DFT_FUNC_NAME)

        # Now implement the function
        block = func.append_basic_block(name="entry")
        builder = ir.IRBuilder(block)
        a, b = func.args
        result = builder.add(a, b, name="res")
        builder.ret(result)

        main_func = ir.Function(module, ir.FunctionType(LLVM_INT_TYPE, ()), name="main")
        builder = ir.IRBuilder(main_func.append_basic_block(name="entry"))
        result = builder.call(
            func, args=[LLVM_INT_TYPE(1), LLVM_INT_TYPE(3)], name="res"
        )
        builder.ret(result)

        return module

    def gen_ir_strlen_module(self):
        fnty = ir.FunctionType(LLVM_INT_TYPE, (LLVM_STR_TYPE,))

        # Create an empty module...
        module = ir.Module(name=__file__)
        module.triple = targets.get_default_triple()

        # declare C lib api, No need to implement it
        strlen_func = ir.Function(
            module, ir.FunctionType(LLVM_INT_TYPE, (LLVM_STR_TYPE,)), name="strlen"
        )

        func = ir.Function(module, fnty, name=DFT_FUNC_NAME)

        # Now implement the function
        block = func.append_basic_block(name="entry")
        builder = ir.IRBuilder(block)
        s = func.args
        strlen_res = builder.call(strlen_func, s)
        builder.ret(strlen_res)
        return module

    def gen_ir_strcat_module(self):
        fnty = ir.FunctionType(LLVM_STR_TYPE, (LLVM_STR_TYPE, LLVM_STR_TYPE))

        # Create an empty module...
        module = ir.Module(name=__file__)
        module.triple = targets.get_default_triple()

        # declare C lib api, No need to implement it
        strconcat_func = ir.Function(
            module,
            ir.FunctionType(LLVM_STR_TYPE, (LLVM_STR_TYPE, LLVM_STR_TYPE)),
            name="str_concat",
        )

        func = ir.Function(module, fnty, name=DFT_FUNC_NAME)

        # Now implement the function
        block = func.append_basic_block(name="entry")
        builder = ir.IRBuilder(block)
        a, b = func.args
        strconcat_res = builder.call(strconcat_func, (a, b))
        builder.ret(strconcat_res)
        return module

    @staticmethod
    def gen_ir_float_module():
        # Now implement the function
        module, block = prepare_basic_block(
            ir.FunctionType(LLVM_FLOAT_TYPE, (LLVM_FLOAT_TYPE, LLVM_FLOAT_TYPE))
        )
        builder = ir.IRBuilder(block)
        a, b = block.function.args
        result = builder.fadd(a, b, name="res")
        builder.ret(result)
        return module

    @staticmethod
    def gen_ir_const_str_list_module():
        elem_type = LLVM_STR_TYPE
        arr_sz = 2
        arr_type = LLVM_LIST_CONTAINER_TYPE(elem_type, arr_sz)

        # Now implement the function
        module, block = prepare_basic_block(
            ir.FunctionType(ir.PointerType(arr_type), ())
        )

        str_var1 = str_value_to_llvm_storage(module, "A", ".str0")
        str_var2 = str_value_to_llvm_storage(module, "B", ".str1")

        arr_var = ir.GlobalVariable(module, arr_type, name=".list0")
        arr_var.global_constant = True
        arr_var.unnamed_addr = True
        arr_var.align = None
        arr_var.linkage = "private"
        indices = [LLVM_INT_TYPE(0), LLVM_INT_TYPE(0)]
        arr_var.initializer = Constant(
            arr_type, [str_var1.gep(indices), str_var2.gep(indices)]
        )

        builder = ir.IRBuilder(block)
        # elem_ptr = builder.gep(arr_var, indices)
        # result = builder.load(elem_ptr)
        # builder.ret(result)
        builder.ret(arr_var)
        return module

    @staticmethod
    def gen_ir_str_list_module():
        elem_type = LLVM_STR_TYPE
        arr_sz = 2
        arr_type = LLVM_LIST_CONTAINER_TYPE(elem_type, arr_sz)

        # Now implement the function
        module, block = prepare_basic_block(
            ir.FunctionType(ir.PointerType(arr_type), ())
        )

        str_var1 = str_value_to_llvm_storage(module, "A", ".str0")
        str_var2 = str_value_to_llvm_storage(module, "B", ".str1")

        arr_var = ir.GlobalVariable(module, arr_type, name=".list0")
        arr_var.global_constant = False
        arr_var.unnamed_addr = True
        arr_var.align = None
        arr_var.linkage = "private"

        arr_var.initializer = None

        indices = [LLVM_INT_TYPE(0), LLVM_INT_TYPE(0)]
        builder = ir.IRBuilder(block)

        ptr1 = builder.gep(arr_var, indices)
        builder.store(str_var1.gep(indices), ptr1)

        ptr2 = builder.gep(arr_var, [LLVM_INT_TYPE(0), LLVM_INT_TYPE(1)])
        builder.store(str_var2.gep(indices), ptr2)
        builder.ret(arr_var)
        return module

    def test_str_len(self):
        c_lib = ctypes.CDLL("/usr/lib/x86_64-linux-gnu/libc.so.6")
        self.assertEqual(c_lib.strlen(b"hello"), 5)

    @staticmethod
    def check_const_list(py_list: list):
        if py_list:
            return not isinstance(py_list[0], GlobalVariable)
        else:
            return True

    @staticmethod
    def build_list_storage(
        module, builder: ir.IRBuilder, elem_type, py_list: list, arr_name
    ):
        is_const = LLVMRunnerTest.check_const_list(py_list)
        arr_sz = len(py_list)
        arr_type = ir.ArrayType(elem_type, arr_sz)
        arr_var = ir.GlobalVariable(module, arr_type, name=arr_name)
        arr_var.global_constant = is_const
        arr_var.unnamed_addr = True
        arr_var.align = None
        arr_var.linkage = "private"

        if is_const:
            arr_var.initializer = ir.Constant(arr_type, [elem_type(x) for x in py_list])
        else:
            for i, global_var in enumerate(py_list):
                ptr = builder.gep(
                    arr_var,
                    [LLVM_INT_TYPE(0), LLVM_INT_TYPE(i)],
                )
                gval = builder.load(global_var)
                builder.store(gval, ptr)

        return arr_var

    @staticmethod
    def build_list_globalvar(
        module, builder: ir.IRBuilder, elem_type, py_list, list_name
    ):
        storage_var = LLVMRunnerTest.build_list_storage(
            module, builder, elem_type, py_list, list_name + ".storage"
        )

        list_type = ir.LiteralStructType([LLVM_INT_TYPE, elem_type.as_pointer()])
        list_var = ir.GlobalVariable(module, list_type, name=list_name)
        list_var.global_constant = True
        list_var.unnamed_addr = True
        list_var.align = None
        list_var.linkage = "private"
        list_var.initializer = ir.Constant(
            list_type,
            [
                LLVM_INT_TYPE(len(py_list)),
                storage_var.gep([LLVM_INT_TYPE(0), LLVM_INT_TYPE(0)]),
            ],
        )
        return list_var

    @staticmethod
    def gen_ir_int_list():
        elem_type = LLVM_INT_TYPE
        module, block = prepare_basic_block(
            ir.FunctionType(get_llvm_list_full_type(elem_type).as_pointer(), ())
        )

        builder = ir.IRBuilder(block)

        list_var = LLVMRunnerTest.build_list_globalvar(
            module, builder, elem_type, [2, 3], ".list0"
        )

        builder.ret(list_var)
        return module

    @staticmethod
    def gen_ir_intlist_list():
        elem_type = get_llvm_list_full_type(LLVM_INT_TYPE)
        module, block = prepare_basic_block(
            ir.FunctionType(get_llvm_list_full_type(elem_type).as_pointer(), ())
        )
        builder = ir.IRBuilder(block)

        list_var1 = LLVMRunnerTest.build_list_globalvar(
            module, builder, LLVM_INT_TYPE, [2, 3], ".list0"
        )
        list_var2 = LLVMRunnerTest.build_list_globalvar(
            module, builder, LLVM_INT_TYPE, [4, 5, 6], ".list1"
        )

        complex_list = LLVMRunnerTest.build_list_globalvar(
            module, builder, elem_type, [list_var1, list_var2], ".complexlist0"
        )
        builder.ret(complex_list)
        return module

    @staticmethod
    def gen_ir_outputvar(var_name):
        module, block = prepare_basic_block(ir.FunctionType(LLVM_INT_TYPE, ()))
        gv = numeric_value_to_llvm_globalvar(module, 10, LLVM_INT_TYPE, var_name)
        builder = ir.IRBuilder(block)
        result = builder.load(gv)
        sum = builder.add(result, LLVM_INT_TYPE(1))
        builder.store(sum, gv)
        builder.ret(result)
        return module

    def test_llvm_list_type(self):
        list_typ = get_llvm_list_full_type(LLVM_INT_TYPE)
        self.assertTrue(isinstance(list_typ, ir.LiteralStructType))
        self.assertTrue(is_llvm_list_type(list_typ))
        self.assertFalse(is_llvm_list_type(LLVM_INT_TYPE))
        self.assertFalse(is_llvm_list_type(LLVM_LIST_CONTAINER_TYPE(LLVM_INT_TYPE, 2)))

    def test_llvm_run(self):
        runner = LLVMRunner()
        self.assertEqual(runner.run(self.gen_ir_i32_module(), 1, 3), 4)
        self.assertEqual(runner.run(self.gen_ir_float_module(), 1.5, 3), 4.5)
        self.assertEqual(runner.run(self.gen_ir_strlen_module(), b"hellome"), 7)
        self.assertEqual(runner.run(self.gen_ir_strlen_module(), b"helloworld1"), 11)
        self.assertEqual(
            runner.run(self.gen_ir_strcat_module(), b"hello", b"world"), b"helloworld"
        )

        res = runner.run(self.gen_ir_const_str_list_module())
        self.assertEqual([x.decode("utf-8") for x in res.contents], ["A", "B"])

        res = runner.run(self.gen_ir_str_list_module())
        self.assertEqual([x.decode("utf-8") for x in res.contents], ["A", "B"])

        res = runner.run(self.gen_ir_int_list())
        self.assertEqual(res.contents.count, 2)
        self.assertEqual(res.contents.elems[: res.contents.count], [2, 3])

        res = runner.run(self.gen_ir_intlist_list())
        self.assertEqual(res.contents.count, 2)
        self.assertEqual(
            [x.elems[: x.count] for x in res.contents.elems[: res.contents.count]],
            [[2, 3], [4, 5, 6]],
        )

        self.assertEqual(runner.run(self.gen_ir_outputvar("a")), 10)
        self.assertEqual(
            c_int.from_address(runner.get_global_value_address("a")).value, 11
        )


if __name__ == "__main__":
    unittest.main()
